#!/bin/bash
# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
#
# Called with following arguments:
# 1 - Path to the downloads folder which is typically
#     ${TENSORFLOW_ROOT}/tensorflow/lite/micro/tools/make/downloads
# 2 - Generated source directory
# 3 - (optional) TENSORFLOW_ROOT: path to root of the TFLM tree (relative to directory from where the script is called).
#
# This script is called from the Makefile and uses the following convention to
# enable determination of sucess/failure:
#
#   - If the script is successful, the only output on stdout should be SUCCESS.
#     The makefile checks for this particular string.
#
#   - Any string on stdout that is not SUCCESS will be shown in the makefile as
#     the cause for the script to have failed.
#
#   - Any other informational prints should be on stderr.

set -e

source ${3}tensorflow/lite/micro/tools/make/bash_helpers.sh

DOWNLOADS_DIR=${1}
GENERATED_SRCS_DIR=${2}

if [ ! -d ${DOWNLOADS_DIR} ]; then
  echo "The top-level downloads directory: ${DOWNLOADS_DIR} does not exist."
  exit 1
fi

# Optimize downloaded model with Vela for Ethos-U.
# See tensorflow/lite/micro/kernels/ethos_u/README.md for more info.
MODEL_DIR=${GENERATED_SRCS_DIR}tensorflow/lite/micro/models
CONVERTED_PERSON_MODEL_INT8=${MODEL_DIR}/person_detect_model_data_vela.cc
PERSON_MODEL_HEADER=${MODEL_DIR}/person_detect_model_data.h

if [ ! -f ${CONVERTED_PERSON_MODEL_INT8} ]; then

  # Install ethos-u-vela if not already installed.
  set +e
  pip show ethos-u-vela >&2
  retVal=$?
  set -e
  if [ $retVal -ne 0 ]; then
    TEMPFILE=$(mktemp -d)/
    python3 -m venv $TEMPFILE
    source $TEMPFILE/bin/activate
    pip install numpy==1.26.4 >&2
    pip install --prefer-binary ethos-u-vela >&2
  fi

  # Compile an optimized .tflite version for Ethos-U.
  vela --accelerator-config=ethos-u55-256 ${DOWNLOADS_DIR}/../../../models/person_detect.tflite \
       --output-dir ${MODEL_DIR} >&2

  if [ $retVal -ne 0 ]; then
    deactivate
  fi

  # Convert .tflite back to C array.
  echo "// This file is generated by $0." > ${CONVERTED_PERSON_MODEL_INT8}
  echo '#include "tensorflow/lite/micro/models/person_detect_model_data.h"' >> \
       ${CONVERTED_PERSON_MODEL_INT8}
  echo -n "const " >> ${CONVERTED_PERSON_MODEL_INT8}
  xxd -i ${MODEL_DIR}/person_detect_vela.tflite >> ${CONVERTED_PERSON_MODEL_INT8}
  sed  -i 's/gen_cortex_m_corstone_300_cortex_m55_.*genfiles_tensorflow_lite_micro_models_person_detect_vela_tflite/g_person_detect_model_data/' \
      ${CONVERTED_PERSON_MODEL_INT8}
  sed -i 's/^const unsigned char g_person_detect_model_data/alignas\(16\) &/' ${CONVERTED_PERSON_MODEL_INT8}
  SIZE=$(sed -E -n -e 's/^.*g_person_detect_model_data_len = ([0-9]+);/\1/p' ${CONVERTED_PERSON_MODEL_INT8})
  sed -i 's/^.*g_person_detect_model_data_len.*$//' ${CONVERTED_PERSON_MODEL_INT8}
  sed -E -i "s/(^constexpr.*g_person_detect_model_data_size = )([0-9]+);/\1$SIZE;/" ${PERSON_MODEL_HEADER}
fi

echo "SUCCESS"
